﻿#pragma once
#include <assert.h>
#include <iostream>
#include <vector>
using namespace std;

template<class T>
class AVL
{
private:
	typedef struct BTNode
	{
		T _data;
		BTNode* _parent = nullptr;
		BTNode* _left = nullptr;
		BTNode* _right = nullptr;
		int _bf;
	}BTNode;

	BTNode* _root;



	// 右单旋
	void RotateR(BTNode* parent)
	{
		BTNode* subL = parent->_left;
		BTNode* subLR = subL->_right;
		parent->_left = subLR;
		subL->_right = parent;
		// subLR可能为空
		if (subLR)
			subLR->_parent = parent;

		// 如果要右旋的节点是根节点，旋上去的节点要记录为新的根节点
		BTNode* parentParent = parent->_parent;
		if(parentParent == nullptr)
		{
			_root = subL;
		}
		else
		{
			if (parentParent->_left == parent)
			{
				parentParent->_left = subL;
			}
			else
			{
				parentParent->_right = subL;
			}
		}
		subL->_parent = parent->_parent;
		parent->_parent = subL;

		// 更新平衡因子
		subL->_bf = 0;
		parent->_bf = 0;
	}

	// 左单旋
	void RotateL(BTNode* parent)
	{
		BTNode* subR = parent->_right;
		BTNode* subRL = subR->_left;
		parent->_right = subRL;
		subR->_left = parent;
		if (subRL)
		{
			subRL->_parent = parent;
		}
		BTNode* parentParent = parent->_parent;
		if (parentParent == nullptr)
		{
			_root = subR;
		}
		else
		{
			if (parentParent->_left == parent)
			{
				parentParent->_left = subR;
			}
			else
			{
				parentParent->_right = subR;
			}
		}
		subR->_parent = parentParent;
		parent->_parent = subR;

		subR->_bf = 0;
		parent->_bf = 0;
	}

	// 左右双旋，处理非纯左子树高的情况
	void RotateLR(BTNode* parent)
	{
		BTNode* subL = parent->_left;
		BTNode* subLR = subL->_right;
		int LRbf = subLR->_bf;
		RotateL(parent->_left);
		RotateR(parent);
		if (LRbf == -1)
		{
			// 			    p
			// 	   subL           (h)
			// (h)     subLR
			//      (h-1)  (h-1)
			//     newNode
			parent->_bf = 1;
		}
		else if (LRbf == 1)
		{
			// 			    p
			// 	   subL           (h)
			// (h)     subLR
			//      (h-1)  (h-1)
			//            newNode
			subL->_bf = -1;
		}
		else if (LRbf == 0)
		{
			//         p
			//   subL
			//      subLR
			// 该情况只需把三个节点的平衡因子都置零，
			// 由于左单旋加右单旋会默认把三个节点的平衡因子置零，
			// 所以不用另外处理
			;
		}
		else
		{
			assert(false);
		}
	}

	// 右左双旋，处理非纯右子树高的情况
	void RotateRL(BTNode* parent)
	{
		BTNode* subR = parent->_right;
		BTNode * subRL = subR->_left;
		int RLbf = subRL->_bf;

		RotateR(parent->_right);
		RotateL(parent);
		if (RLbf == -1)
		{
			//	         p
			//   (h)	       subR           
			//          subRL        (h)
			//     (h-1)   (h-1)
			//    newNode
			subR->_bf = 1;
		}
		else if (RLbf == 1)
		{
			//	         p
			//   (h)	       subR           
			//          subRL        (h)
			//      (h-1)   (h-1)
			//             newNode
			parent->_bf = -1;
		}
		else if (RLbf == 0)
		{
			//          p
			//            subR
			//        subRL
			// 该情况只需把三个节点的平衡因子都置零，
			// 由于左单旋加右单旋会默认把三个节点的平衡因子置零，
			// 所以不用另外处理
			;
		}
		else
		{
			assert(false);
		}
	}

public:

	AVL() : _root(nullptr)
	{}

	bool insert(T val)
	{
		BTNode* tmp = new BTNode;
		if (tmp == nullptr)
		{
			perror("insert error!\n");
			return false;
		}
		tmp->_data = val;
		tmp->_bf = 0;

		// 如果当前是空树
		if (_root == nullptr)
		{
			_root = tmp;
			return true;
		}
		BTNode* parent = nullptr;
		BTNode* cur = _root;

		// 找到插入的位置
		while (cur != nullptr)
		{
			if (val < cur->_data)
			{
				parent = cur;
				cur = cur->_left;
			}
			else if(val > cur->_data)
			{
				parent = cur;
				cur = cur->_right;
			}
			else
			{
				return false;
			}
		}

		// 进行插入
		if (cur == parent->_left)
		{
			parent->_left = tmp;
			tmp->_parent = parent;
			cur = parent->_left;
		}
		else if (cur == parent->_right)
		{
			parent->_right = tmp;
			tmp->_parent = parent;
			cur = parent->_right;
		}

		// 更新平衡因子
		while (parent)
		{
			if (cur == parent->_left)
			{
				parent->_bf--;
			}
			else
			{
				parent->_bf++;
			}

			if (parent->_bf == 0)
			{
				break;
			}
			else if (parent->_bf == 1 || parent->_bf == -1)
			{
				cur = parent;
				parent = parent->_parent;
			}
			else if (parent->_bf == 2 || parent->_bf == -2)
			{
				// 旋转
				if (parent->_bf == -2)
				{
					if (parent->_left->_bf == -1)
					{
						RotateR(parent);
					}
					else if (parent->_left->_bf == 1)
					{
						RotateLR(parent);
					}
					else
					{
						assert(false);
					}
				}
				else if (parent->_bf == 2)
				{
					if (parent->_right->_bf == 1)
					{
						RotateL(parent);
					}
					else if (parent->_right->_bf == -1)
					{
						RotateRL(parent);
					}
					else
					{
						assert(false);
					}
				}
				break;
			}
			else
			{
				assert(false);
			}
		}

		return true;
	}

	bool find(const T& val)
	{
		BTNode* cur = _root;
		while (cur != nullptr)
		{
			if (val < cur->_data)
			{
				cur = cur->_left;
			}
			else if (val > cur->_data)
			{
				cur = cur->_right;
			}
			else
			{
				return true;
			}
		}
		return false;
	}


	// AVL树平衡检测
	int _Height(BTNode* root)
	{
		if (root == nullptr)
			return 0;
		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);
		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}

	bool _IsBalanceTree(BTNode* root)
	{
		// 空树也是AVL树
		if (nullptr == root)
			return true;
		// 计算pRoot结点的平衡因⼦：即pRoot左右⼦树的⾼度差
		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);
		int diff = rightHeight - leftHeight;
		// 如果计算出的平衡因⼦与pRoot的平衡因⼦不相等，或者
		// pRoot平衡因⼦的绝对值超过1，则⼀定不是AVL树
		if (abs(diff) >= 2)
		{
			cout << root->_data << "height error" << endl;
			return false;
		}
		if (root->_bf != diff)
		{
			cout << root->_data << "bf error" << endl;
			return false;
		}
		// pRoot的左和右如果都是AVL树，则该树⼀定是AVL树
		return _IsBalanceTree(root->_left) && _IsBalanceTree(root->_right);
	}

	bool IsBalanceTree()
	{
		return _IsBalanceTree(_root);
	}

	int Height()
	{
		return _Height(_root);
	}

	size_t _size(BTNode* root)
	{
		if (!root)
			return 0;
		size_t cnt1 = _size(root->_left);
		size_t ret = 1;
		size_t cnt2 = _size(root->_right);
		return cnt1 + cnt2 + ret;
	}

	size_t size()
	{
		return _size(_root);
	}
};
